# tts/tts_audio.py
# 音频播放和队列管理模块

import queue
import threading
import pyaudio
from typing import Optional


class AudioPlayer:
    """音频播放器，负责音频队列管理和播放"""
    
    def __init__(self, sample_rate: int):
        self.sample_rate = sample_rate
        self.audio_queue = queue.Queue()
        self.player_thread: Optional[threading.Thread] = None
        self._stop_event = threading.Event()
        self._pyaudio_instance: Optional[pyaudio.PyAudio] = None
        self._stream: Optional[pyaudio.Stream] = None
    
    def start(self):
        """启动音频播放器线程"""
        if not self.player_thread or not self.player_thread.is_alive():
            self.player_thread = threading.Thread(target=self._audio_player_thread)
            self.player_thread.daemon = True
            self.player_thread.start()
            print("\n[TTS] 音频播放器服务已启动。")
    
    def stop(self):
        """停止音频播放器"""
        self._stop_event.set()
        if self.player_thread and self.player_thread.is_alive():
            self.player_thread.join(timeout=2)
        self._cleanup()
    
    def stop_immediately(self):
        """立即停止音频播放器（用于中断）"""
        # 立即清空队列
        self.clear_queue()
        # 立即停止音频流，但不设置停止事件（保持线程运行）
        if self._stream:
            try:
                self._stream.stop_stream()
            except:
                pass
    
    def put_audio_chunk(self, chunk: bytes):
        """向音频队列添加音频块"""
        if chunk is not None and len(chunk) > 0:
            self.audio_queue.put(chunk)
    
    def clear_queue(self):
        """清空音频队列"""
        while not self.audio_queue.empty():
            try: 
                self.audio_queue.get_nowait()
            except queue.Empty: 
                break
    
    def _audio_player_thread(self):
        """音频播放器线程主循环"""
        try:
            self._pyaudio_instance = pyaudio.PyAudio()
            self._stream = self._pyaudio_instance.open(
                format=pyaudio.paInt16, 
                channels=1, 
                rate=self.sample_rate, 
                output=True,
                frames_per_buffer=1024  # 增加缓冲区大小以改善音频质量
            )
            
            # 预启动音频流，确保第一个chunk能立即播放
            self._stream.start_stream()
            
            while not self._stop_event.is_set():
                try:
                    # 使用超时避免无限阻塞
                    chunk = self.audio_queue.get(timeout=0.1)
                    if chunk is None: 
                        continue 
                    
                    # 确保流处于活动状态，如果被中断则重新启动
                    if not self._stream.is_active(): 
                        try:
                            self._stream.start_stream()
                        except Exception as e:
                            # 如果启动失败，尝试重新创建流
                            try:
                                self._stream.close()
                                self._stream = self._pyaudio_instance.open(
                                    format=pyaudio.paInt16, 
                                    channels=1, 
                                    rate=self.sample_rate, 
                                    output=True,
                                    frames_per_buffer=1024
                                )
                                self._stream.start_stream()
                            except Exception as e2:
                                print(f"\n[TTS] 音频流重建失败: {e2}")
                                continue
                    
                    # 写入音频数据
                    self._stream.write(chunk)
                except queue.Empty:
                    continue
                except Exception as e:
                    if not self._stop_event.is_set():
                        print(f"\n[TTS] 音频播放错误: {e}")
                    break
        finally:
            self._cleanup()
            print("\n[TTS] 音频播放器服务已关闭。")
    
    def _cleanup(self):
        """清理音频资源"""
        if self._stream:
            try:
                self._stream.stop_stream()
                self._stream.close()
            except:
                pass
            self._stream = None
        
        if self._pyaudio_instance:
            try:
                self._pyaudio_instance.terminate()
            except:
                pass
            self._pyaudio_instance = None


class AudioQueueManager:
    """音频队列管理器，提供更高级的队列操作"""
    
    def __init__(self, audio_player: AudioPlayer):
        self.audio_player = audio_player
        self._generation_id = 0
        self._generation_lock = threading.Lock()
    
    def get_current_generation_id(self) -> int:
        """获取当前代际ID"""
        with self._generation_lock:
            return self._generation_id
    
    def increment_generation(self) -> int:
        """增加代际ID并返回新的ID"""
        with self._generation_lock:
            self._generation_id += 1
            return self._generation_id
    
    def is_generation_current(self, gen_id: int) -> bool:
        """检查代际ID是否仍然有效"""
        with self._generation_lock:
            return gen_id == self._generation_id
    
    def interrupt_and_prepare_new_task(self) -> int:
        """中断当前任务并准备新任务"""
        print("\n[系统] 检测到输入，正在中断上一轮对话...")
        
        # 立即停止音频播放
        self.audio_player.stop_immediately()
        
        # 清空音频队列
        self.audio_player.clear_queue()
        
        # 增加代际ID
        new_gen_id = self.increment_generation()
        
        return new_gen_id
